MODULE md1_density

   USE md1_global_consts
   USE md1_global_types
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: md1_density_init,        &
             md1_density_free,        &
             md1_get_rho_on_grid

   REAL(REALK), SAVE :: md1_rho_thresh = 1d-10
   REAL(REALK), ALLOCATABLE, TARGET, SAVE :: md1_rho_on_grid(:,:,:)
   LOGICAL, SAVE :: md1_density_built_flag = .FALSE.

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE md1_density_init(grid)

      IMPLICIT NONE  
      TYPE(md1_FE_grid), INTENT(IN) :: grid
      ALLOCATE ( md1_rho_on_grid(grid%n(1),grid%n(2),grid%n(3)) )

   END SUBROUTINE md1_density_init

!-------------------------------------------------------------------------------

   SUBROUTINE md1_density_free

      IMPLICIT NONE  
      IF (ALLOCATED(md1_rho_on_grid)) DEALLOCATE (md1_rho_on_grid)
      md1_density_built_flag = .FALSE.

   END SUBROUTINE md1_density_free

!-------------------------------------------------------------------------------

   SUBROUTINE md1_get_rho_on_grid(cMO,nbas,nocc,grid,dens)

      USE md1_basis_set, ONLY: md1_get_basis_set

      IMPLICIT NONE  
      REAL(REALK),       INTENT(IN) :: cMO(nbas,nocc)
      INTEGER(INTK),     INTENT(IN) :: nbas, nocc
      TYPE(md1_FE_grid), INTENT(IN) :: grid
      REAL(REALK),       POINTER    :: dens(:,:,:)

      TYPE(md1_basis), POINTER :: basis

      IF (md1_density_built_flag .AND. ALLOCATED(md1_rho_on_grid)) THEN
         !grid%f(:,:,:) = md1_rho_on_grid(:,:,:)
         dens => md1_rho_on_grid(:,:,:)
      ELSE
         CALL md1_get_basis_set(basis)
         CALL md1_build_rho_on_grid(basis,cMO,nocc,grid)
         md1_density_built_flag = .TRUE.
         !grid%f(:,:,:) = md1_rho_on_grid(:,:,:)
         dens => md1_rho_on_grid(:,:,:)
      END IF

   END SUBROUTINE md1_get_rho_on_grid

!-------------------------------------------------------------------------------

   SUBROUTINE md1_build_rho_on_grid(basis,cMO,nocc,grid)

      IMPLICIT NONE  
      TYPE(md1_basis),   INTENT(IN) :: basis
      REAL(REALK),       INTENT(IN) :: cMO(:,:)
      INTEGER(INTK),     INTENT(IN) :: nocc
      TYPE(md1_FE_grid), INTENT(IN) :: grid

      REAL(REALK), POINTER :: rho(:,:,:)
      REAL(REALK)   :: MO(grid%n(1),grid%n(2),grid%n(3),nocc)
      REAL(REALK)   :: AO(grid%n(1),grid%n(2),grid%n(3))
      REAL(REALK)   :: sphAO(grid%n(1),grid%n(2),grid%n(3))
      REAL(REALK)   :: tmp, thresh
      INTEGER(INTK) :: iMO, Ish, L,M, p
      LOGICAL :: empty

      MO = zero   

      ! Project MOs onto grid
      shells: DO Ish = 1, basis%nshell

         CALL md1_project_spherical_AO(basis,Ish,grid,sphAO,empty)
         IF (empty) CYCLE shells

         L = basis%L(Ish)
         p = basis%car_AO(ISh)

         AOs: DO M = 1, basis%LM(L)%car_nlm

            CALL md1_project_full_AO(basis,Ish,M,grid,sphAO,AO)
            tmp = MAXVAL(ABS(AO))

            MOs: DO iMO = 1, nocc
               IF (tmp * ABS(cMO(p,iMO)) < md1_rho_thresh) CYCLE
               MO(:,:,:,iMO) = MO(:,:,:,iMO) + cMO(p,iMO) * AO(:,:,:)
            END DO MOs

            p = p+1
         END DO AOs

      END DO shells

      ! Build rho from projected MOs
      rho => md1_rho_on_grid(:,:,:)
      rho(:,:,:) = zero
      DO iMO = 1, nocc
         IF ( MAXVAL(ABS(MO(:,:,:,iMO))) < md1_rho_thresh ) CYCLE
         rho(:,:,:) = rho(:,:,:) + MO(:,:,:,iMO) * MO(:,:,:,iMO)
      END DO

      ! Doubly occupied MOs
      rho(:,:,:) = two*rho(:,:,:)

   END SUBROUTINE md1_build_rho_on_grid

!-------------------------------------------------------------------------------

   SUBROUTINE md1_project_spherical_AO(basis,Ish,grid,sphAO,empty)

      IMPLICIT NONE
      TYPE(md1_basis),   INTENT(IN)  :: basis
      INTEGER(INTK),     INTENT(IN)  :: Ish
      TYPE(md1_FE_grid), INTENT(IN)  :: grid
      REAL(REALK),       INTENT(OUT) :: sphAO(:,:,:)
      LOGICAL,           INTENT(OUT) :: empty

      REAL(REALK)   :: tmp(grid%n(1),grid%n(2),grid%n(3))
      REAL(REALK)   :: expnx(grid%n(1)), expny(grid%n(2)), expnz(grid%n(3))
      REAL(REALK)   :: sh_cntr(3), gpt, X,Y,Z
      INTEGER(INTK) :: i,j,k, p,ip1,ip2

      empty = .TRUE.
      sh_cntr(:) = basis%xyz(:,basis%atom(Ish))

      sphAO(:,:,:) = zero
      ip1 = basis%exp_id(Ish)
      ip2 = ip1 + basis%nprim(Ish) - 1

      primitives: DO p = ip1, ip2

         ! x
         DO i = 1, grid%n(1)
            gpt = (i-1)*grid%h(1) + grid%lo(1)
            X = gpt - sh_cntr(1)
            expnx(i) = EXP(-basis%exp(p) * X*X)
         END DO
         IF (MAXVAL(expnx(:)) < md1_rho_thresh) CYCLE primitives
         ! y
         DO i = 1, grid%n(2)
            gpt = (i-1)*grid%h(2) + grid%lo(2)
            Y = gpt - sh_cntr(2)
            expny(i) = EXP(-basis%exp(p) * Y*Y)
         END DO
         IF (MAXVAL(expny(:)) < md1_rho_thresh) CYCLE primitives
         ! z
         DO i = 1, grid%n(3)
            gpt = (i-1)*grid%h(3) + grid%lo(3)
            Z = gpt - sh_cntr(3)
            expnz(i) = EXP(-basis%exp(p) * Z*Z)
         END DO
         IF (MAXVAL(expnz(:)) < md1_rho_thresh) CYCLE primitives

         ! xyz product
         DO k = 1, grid%n(3)
            DO j = 1, grid%n(2)
               DO i = 1, grid%n(1)
                  tmp(i,j,k) = expnx(i)*expny(j)*expnz(k)
               END DO
            END DO
         END DO

         IF (MAXVAL(tmp) < md1_rho_thresh) CYCLE primitives
         tmp = basis%C(p) * tmp
         IF (MAXVAL(tmp) < md1_rho_thresh) CYCLE primitives

         sphAO = sphAO + tmp
         empty = .FALSE.

      END DO primitives

   END SUBROUTINE md1_project_spherical_AO

!-------------------------------------------------------------------------------

   SUBROUTINE md1_project_full_AO(basis,Ish,M,grid,sphAO,AO)

      IMPLICIT NONE
      TYPE(md1_basis),   INTENT(IN)  :: basis
      INTEGER(INTK),     INTENT(IN)  :: Ish, M
      TYPE(md1_FE_grid), INTENT(IN)  :: grid
      REAL(REALK),       INTENT(IN)  :: sphAO(:,:,:)
      REAL(REALK),       INTENT(OUT) :: AO(:,:,:)

      REAL(REALK)   :: facX(grid%n(1)), facY(grid%n(2)), facZ(grid%n(3))
      REAL(REALK)   :: sh_cntr(3)
      REAL(REALK)   :: gpt, X,Y,Z
      INTEGER(INTK) :: i,j,k, L

      sh_cntr(:) = basis%xyz(:,basis%atom(Ish))

      L = basis%L(Ish)

      ! x
      DO i = 1, grid%n(1)
         gpt = (i-1)*grid%h(1) + grid%lo(1)
         X = gpt - sh_cntr(1)
         facX(i) = X**basis%LM(L)%tuv(1,M)
      END DO
      ! y
      DO i = 1, grid%n(2)
         gpt = (i-1)*grid%h(2) + grid%lo(2)
         Y = gpt - sh_cntr(2)
         facY(i) = Y**basis%LM(L)%tuv(2,M)
      END DO
      ! z
      DO i = 1, grid%n(3)
         gpt = (i-1)*grid%h(3) + grid%lo(3)
         Z = gpt - sh_cntr(3)
         facZ(i) = Z**basis%LM(L)%tuv(3,M)
      END DO

      ! x,y,z product
      DO k = 1, grid%n(3)
         DO j = 1, grid%n(2)
            DO i = 1, grid%n(1)
               AO(i,j,k) = facX(i)*facY(j)*facZ(k) * sphAO(i,j,k)
            END DO
         END DO
      END DO

   END SUBROUTINE md1_project_full_AO

!-------------------------------------------------------------------------------

END MODULE md1_density
